import os
from typing import List

from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers.bm25 import BM25Retriever
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from pydantic import BaseModel, Field

from langchain_workflow.embedding import embedding_model
from logger import logger


class FaissConfig(BaseModel):
    persist_dir: str = Field(default="vectordb/faiss/documents")
    bm25_retriever_params: dict = Field(default={"k": 5})
    dense_retriever_params: dict = Field(default={"search_kwargs": {"k": 10}})
    ensemble_retriever_params: dict = Field(default={"weights": [0.5, 0.5]})


class FaissHybridRetrieverFactory:
    def __init__(self, documents: List[Document], config: FaissConfig) -> None:
        self.documents = self.__post_documents(documents)
        self.config = config

        self.embedding_func = embedding_model
        self.vectordb = self._init_or_load_faiss()
        self.bm25_retriever = self._init_bm25_retriever()
        self.faiss_retriever = self._init_faiss_retriever()
        # 延迟初始化 ensemble_retriever
        self._ensemble_retriever = None

    def __post_documents(self, documents: List[Document]) -> List[Document]:
        """对文档进行后处理，将元数据中的上下文信息添加到文档内容中

        此函数遍历输入的文档列表，检查每个文档的元数据是否包含"context"字段。
        如果存在，则将上下文信息与原始文档内容合并，以增强文档的信息量。

        Args:
            documents (List[Document]): 待处理的文档列表

        Returns:
            List[Document]: 处理后的文档列表，每个文档可能包含增强的内容
        """
        new_documents = []
        for document in documents:
            if document.metadata.get("context"):
                context = "||".join(document.metadata.values())
                document.page_content = (
                    f"context: {context} \n content: {document.page_content}"
                )

            new_documents.append(document)

        return new_documents

    def _init_or_load_faiss(self):
        if os.path.exists(self.config.persist_dir):
            logger.info(f"Loading existing FAISS index from {self.config.persist_dir}")
            return FAISS.load_local(
                self.config.persist_dir,
                self.embedding_func,
                allow_dangerous_deserialization=True,
            )
        else:
            logger.info(
                f"Creating new FAISS index and saving to {self.config.persist_dir}"
            )
            vectordb = FAISS.from_documents(
                documents=self.documents,
                embedding=self.embedding_func,
            )
            vectordb.save_local(self.config.persist_dir)
            return vectordb

    def _init_faiss_retriever(self):
        return self.vectordb.as_retriever(
            search_kwargs=self.config.dense_retriever_params
        )

    def _init_bm25_retriever(self) -> BM25Retriever:
        bm25_retriever = BM25Retriever.from_documents(
            self.documents,
            **self.config.bm25_retriever_params,
        )
        return bm25_retriever

    def _init_ensemble_retriever(self):
        ensemble_retriever = EnsembleRetriever(
            retrievers=[self.bm25_retriever, self.faiss_retriever],
            **self.config.ensemble_retriever_params,
        )
        return ensemble_retriever

    @property
    def ensemble_retriever(self):
        if self._ensemble_retriever is None:
            self._ensemble_retriever = self._init_ensemble_retriever()
        return self._ensemble_retriever

    def get_retriever(self) -> EnsembleRetriever:
        """
        获取混合检索器
            topk (int): 检索时返回的最大结果数

        返回:
            EnsembleRetriever: 混合检索器实例
        """
        return self.ensemble_retriever


def get_faiss_hybrid_retriever(documents: List[Document]):
    config = FaissConfig()
    return FaissHybridRetrieverFactory(documents, config).get_retriever()


if __name__ == "__main__":
    from langchain_workflow.splitter import get_all_documents

    all_documents = get_all_documents()

    from langchain_workflow.process.add_context_procesor import ContextSpliter

    all_documents_with_context = ContextSpliter(all_documents).run()
    retriever = get_faiss_hybrid_retriever(all_documents_with_context)
    res = retriever.invoke(
        "2-6PACK Projection Flashlight Interactive Toy Night Lamp For Kids Infant"
    )
    print(res)
